<?php

/**
 * This file is part of Webman AI.
 *
 * @author    walkor<walkor@workerman.net>
 * @copyright walkor<walkor@workerman.net>
 * @link      http://www.workerman.net/
 */

namespace plugin\ai\app\controller;

use Exception;
use plugin\ai\api\Task;
use plugin\ai\app\event\data\EventData;
use plugin\ai\app\event\data\ModelRequestData;
use plugin\ai\app\event\data\ModelResponseData;
use plugin\ai\app\model\AiMessage;
use plugin\ai\app\service\Common;
use plugin\ai\app\service\User;
use plugin\ai\app\model\AiImage;
use support\exception\BusinessException;
use support\Log;
use support\Request;
use support\Response;
use Throwable;
use Webman\Event\Event;
use Webman\Push\Api;
use Workerman\Protocols\Http\Chunk;

/**
 * midjourney画图模块
 */
class TaskController extends Base
{

    /**
     * 不需要登录的方法
     *
     * @var string[]
     */
    protected $noNeedLogin = ['imagine', 'notify', 'status', 'action', 'settings'];

    /**
     * 作图
     * @param Request $request
     * @return Response
     * @throws Exception
     */
    public function imagine(Request $request): ?Response
    {
        $imageData = $request->post('image');
        $prompt = $imageData['prompt'] ?? '';
        if (!$prompt || !is_string($prompt) || !trim($prompt)) {
            return json(['error' => ['message' => "请输入图片描述"]]);
        }
        $prompt = $this->formatPrompt($imageData, $prompt);
        $userId = session('user.id') ?? session('user.uid');
        $model = 'midjourney';
        if (empty($imageData['fast']) && Common::getModelInfo('midjourney-relax')) {
            $model = 'midjourney-relax';
        }

        $modelType = Common::getModelType($model);
        $beforeBalance = $userId ? User::getBalance($userId, $modelType) : 0;
        $useUserBalance = false;
        if ($error = static::tryReduceBalance($model, $useUserBalance)) {
            return json(['error' => ['message' => $error]]);
        }
        if ($userId) {
            $afterBalance = User::getBalance($userId, $modelType);
            Log::info("Midjourney reduceBalance model:$model modelType:$modelType type:imagine userId:$userId balance:$beforeBalance->$afterBalance prompt:$prompt");
        }

        $sessionId = $request->sessionId();
        $remoteIp = $request->getRealIp();
        $connection = $request->connection;
        $modelRequestData = new ModelRequestData();
        $modelRequestData->data = [
            'model' => $model,
            'images' => $imageData['refs'] ?? [],
            'notifyUrl' => $request->header('X-Forwarded-Proto', 'http') . '://' . $request->host() . '/app/ai/task/notify',
            'prompt' => $prompt,
            'allowFast' => $userId && User::isVip($userId),
        ];
        $modelRequestData->options = [
            'complete' => function($json) use ($connection, $userId, $sessionId, $remoteIp, $model, $modelRequestData, $imageData, $prompt, $useUserBalance) {
                $responseData = new ModelResponseData();
                $responseData->data = $json;
                $responseData->modelRequestData = $modelRequestData;
                Event::dispatch('ai.task.imagine.response', $responseData);
                $content = json_encode($responseData->data, JSON_UNESCAPED_UNICODE | JSON_PRETTY_PRINT);
                $connection->send(new Chunk($content));
                $connection->send(new Chunk(''));

                if (isset($json['error']) || ($json['code'] ?? 0)) {
                    $modelType = Common::getModelType($model);
                    if ($useUserBalance) {
                        $beforeBalance = User::getBalance($userId, $modelType);
                        User::addBalance($userId, $modelType);
                        $afterBalance = User::getBalance($userId, $modelType);
                        Log::error("Midjourney Error try addBalance model:$model modelType:$modelType type:imagine userId:$userId balance:$beforeBalance->$afterBalance prompt:$prompt error:$content");
                    } else {
                        Log::error("Midjourney Error type:imagine model:$model modelType:$modelType userId:$userId error:$content");
                    }
                } else {
                    $taskId = $json['taskId'];
                    $image = new AiImage();
                    $image->user_id = $userId;
                    $image->session_id = $sessionId;
                    $image->task_id = $taskId;
                    $image->model = $model;
                    $image->prompt = $imageData['text'] ?? '';
                    $image->image = json_encode($imageData, JSON_UNESCAPED_UNICODE | JSON_PRETTY_PRINT);
                    $image->ip = $remoteIp;
                    $image->type = 'IMAGINE';
                    $image->save();
                }
            }
        ];
        Event::dispatch('ai.task.imagine.request', $modelRequestData);
        Task::imagine($modelRequestData->data, $modelRequestData->options);

        // 向浏览器发送头部响应
        return response("\n")->withHeaders([
            "Content-Type" => "application/json",
            "Transfer-Encoding" => "chunked",
        ]);
    }

    /**
     * 图片操作
     * @param Request $request
     * @return Response|null
     * @throws BusinessException
     */
    public function action(Request $request): ?Response
    {
        $imageData = $request->post('image');
        $taskId = $request->post('taskId');
        $customId = $request->post('customId');
        $userId = session('user.id') ?? session('user.uid');
        $sessionId = $request->sessionId();
        $remoteIp = $request->getRealIp();

        $image = AiImage::where('task_id', $taskId)->first();
        if (!$image) {
            return json(['error' => ['message' => '任务不存在']]);
        }
        $model = $image->model;

        $prompt = $imageData['prompt'] ?? '';
        $needReduceBalance = !strpos($customId, '::upsample::');
        $modelType = Common::getModelType($model);
        $beforeBalance = $userId ? User::getBalance($userId, $modelType) : 0;
        $useUserBalance = false;
        if ($needReduceBalance && $error = static::tryReduceBalance($model, $useUserBalance)) {
            return json(['error' => ['message' => $error]]);
        }
        $prompt = $this->formatPrompt($imageData, $prompt);
        if ($needReduceBalance) {
            $afterBalance = $userId ? User::getBalance($userId, $modelType) : 0;
            Log::info("Midjourney reduceBalance model:$model modelType:$modelType type:action customId:$customId taskId:$taskId userId:$userId balance:$beforeBalance->$afterBalance prompt:$prompt");
        }

        $connection = $request->connection;
        $modelRequestData = new ModelRequestData();
        $modelRequestData->data = [
            'model' => $model,
            'taskId' => $taskId,
            'customId' => $customId,
            'prompt' => $prompt,
            'mask' => $request->post('mask'),
            'allowFast' => $userId && User::isVip($userId),
            'notifyUrl' => $request->header('X-Forwarded-Proto', 'http') . '://' . $request->host() . '/app/ai/task/notify',
        ];
        $modelRequestData->options = [
            'complete' => function($json) use ($connection, $modelRequestData, $userId, $sessionId, $imageData, $remoteIp, $model, $needReduceBalance, $taskId, $customId) {
                $responseData = new ModelResponseData();
                $responseData->data = $json;
                $responseData->modelRequestData = $modelRequestData;
                Event::dispatch('ai.task.action.response', $responseData);

                $content = json_encode($responseData->data, JSON_UNESCAPED_UNICODE | JSON_PRETTY_PRINT);
                $connection->send(new Chunk($content));
                $connection->send(new Chunk(''));

                if (isset($json['error']) || ($json['code'] ?? 0)) {
                    $modelType = Common::getModelType($model);
                    if ($needReduceBalance && $userId && User::getBalance($userId, $modelType) > 0) {
                        $beforeBalance = User::getBalance($userId, $modelType);
                        User::addBalance($userId, $modelType);
                        $afterBalance = User::getBalance($userId, $modelType);
                        Log::error("Midjourney Error try addBalance model:$model type:action customId:$customId taskId:$taskId userId:$userId balance:$beforeBalance->$afterBalance error:$content");
                    } else {
                        Log::error("Midjourney Error model:$model type:action customId:$customId taskId:$taskId userId:$userId error:$content");
                    }
                } else {
                    $taskId = $json['taskId'];
                    $image = new AiImage();
                    $image->user_id = $userId;
                    $image->session_id = $sessionId;
                    $image->task_id = $taskId;
                    $image->model = $model;
                    $image->prompt = $imageData['text'] ?? '';
                    $image->image = json_encode($imageData, JSON_UNESCAPED_UNICODE | JSON_PRETTY_PRINT);
                    $image->ip = $remoteIp;
                    $image->save();
                }
            }
        ];
        Event::dispatch('ai.task.action.request', $modelRequestData);
        Task::action($modelRequestData->data, $modelRequestData->options);

        // 向浏览器发送头部响应
        return response("\n")->withHeaders([
            "Content-Type" => "application/json",
            "Transfer-Encoding" => "chunked",
        ]);
    }

    /**
     * 任务状态
     * @param Request $request
     * @return Response|null
     */
    public function status(Request $request): ?Response
    {
        $taskId = $request->get('taskId');
        $image = AiImage::where('task_id', $taskId)->first();
        if (!$image) {
            return json(['error' => ['message' => '任务不存在']]);
        }
        $model = $image->model;
        $connection = $request->connection;
        $modelRequestData = new ModelRequestData();
        $modelRequestData->data = [
            'model' => $model,
            'taskId' => $taskId
        ];
        $modelRequestData->options = [
            'complete' => function($json) use ($connection, $modelRequestData, $taskId, $model, $image) {
                $responseData = new ModelResponseData();
                $responseData->data = $json;
                $responseData->modelRequestData = $modelRequestData;
                Event::dispatch('ai.task.status.response', $responseData);
                $progress = $responseData->data['data']['progress'] ?? '';
                $url = $responseData->data['data']['imageUrl'] ?? '';
                $smallUrl = $responseData->data['data']['smallUrl'] ?? '';
                $thumbUrl = $responseData->data['data']['thumbUrl'] ?? '';
                $content = json_encode($responseData->data, JSON_UNESCAPED_UNICODE | JSON_PRETTY_PRINT);
                if ($progress === '100%' && $url && !$image->image_url) {
                    $image->image_url = $image->image_url ?: $url;
                    $image->small_url = $image->small_url ?: $smallUrl;
                    $image->thumb_url = $image->thumb_url ?: $thumbUrl;
                    $image->data = json_encode($responseData->data['data'], JSON_UNESCAPED_UNICODE | JSON_PRETTY_PRINT);
                    $image->type = $responseData->data['data']['action'];
                    $image->save();
                }
                if (isset($responseData->data['error']) || ($responseData->data['code'] ?? 0) || ($responseData->data['data']['failReason'] ?? false)) {
                    $userId = $image->user_id;
                    $imageData = json_decode($image->image, true);
                    $customId = $imageData['button']['custom_id'] ?? $responseData->data['data']['params']['customId'] ?? '';
                    $needAddBalance = !strpos($customId, '::upsample::') && !$image->data;
                    $modelType = Common::getModelType($image->model);
                    if (!$image->data) {
                        $image->data = json_encode($responseData->data['data'] ?? $responseData->data, JSON_UNESCAPED_UNICODE | JSON_PRETTY_PRINT);
                        $image->save();
                    }
                    if ($needAddBalance && $userId && User::getBalance($userId, $modelType) > 0) {
                        $beforeBalance = User::getBalance($userId, $modelType);
                        User::addBalance($userId, $modelType);
                        $afterBalance = User::getBalance($userId, $modelType);
                        Log::error("Midjourney Error try addBalance modelType:$modelType model:$model type:status customId:$customId taskId:$taskId userId:$userId balance:$beforeBalance->$afterBalance error:$content");
                    } else {
                        Log::error("Midjourney Error modelType:$modelType model:$model type:status customId:$customId taskId:$taskId userId:$userId error:$content");
                    }
                }
                $connection->send(new Chunk($content));
                $connection->send(new Chunk(''));
            }
        ];

        Event::dispatch('ai.task.status.request', $modelRequestData);
        Task::status($modelRequestData->data, $modelRequestData->options);

        // 向浏览器发送头部响应
        return response("\n")->withHeaders([
            "Content-Type" => "application/json",
            "Transfer-Encoding" => "chunked",
        ]);
    }

    /**
     * 接收midjourney_porxy代理发来的通知
     * @param Request $request
     * @return Response
     * @throws BusinessException
     */
    public function notify(Request $request): Response
    {
        $eventData = new EventData($request->post());
        Event::dispatch('ai.task.notify', $eventData);
        $data = $eventData->data;
        $taskId = $data['id'] ?? '';
        $image = AiImage::where('task_id', $taskId)->first();
        if (!$image) {
            return response(json_encode([
                'code' => 0,
                'msg' => 'ok'
            ]));
        }
        $userId = $image->user_id;
        $model = $image->model;
        $progress = $data['progress'] ?? '';
        $url = $data['imageUrl'] ?? $data['image'] ?? '';
        $smallUrl = $data['smallUrl'] ?? '';
        $thumbUrl = $data['thumbUrl'] ?? '';
        if ($progress === '100%' && $url) {
            if (!$image->image_url) {
                $image->image_url = $image->image_url ?: $url;
                $image->small_url = $image->small_url ?: $smallUrl;
                $image->thumb_url = $image->thumb_url ?: $thumbUrl;
                $image->data = json_encode($data, JSON_UNESCAPED_UNICODE | JSON_PRETTY_PRINT);
                $image->type = $data['action'];
                $image->save();
            }
        }
        if ($data['failReason'] ?? '') {
            $needAddBalance = !empty($data['action']) && $data['action'] !== 'UPSCALE' && !$image->data;
            $modelType = Common::getModelType($image->model);
            if (!$image->data) {
                $image->data = json_encode($data, JSON_UNESCAPED_UNICODE | JSON_PRETTY_PRINT);
            }
            $content = json_encode($eventData->data, JSON_UNESCAPED_UNICODE | JSON_PRETTY_PRINT);
            if ($needAddBalance && $userId && User::getBalance($userId, $modelType) > 0) {
                $beforeBalance = User::getBalance($userId, $modelType);
                User::addBalance($userId, $modelType);
                $afterBalance = User::getBalance($userId, $modelType);
                Log::error("Midjourney Error try addBalance modelType:$modelType model:$model type:notify taskId:$taskId userId:$userId balance:$beforeBalance->$afterBalance error:$content");
            } else {
                Log::error("Midjourney Error modelType:$modelType model:$model type:notify taskId:$taskId userId:$userId error:$content");
            }
            $image->type = $image->type ?: $data['action'];
            $image->save();
        }

        try {
            $api = new Api(
                'http://127.0.0.1:3232',
                config('plugin.webman.push.app.app_key'),
                config('plugin.webman.push.app.app_secret')
            );
            $api->trigger($data['state'], 'mj-state-change', $data);
        } catch (Throwable $e) {}
        return response(json_encode([
            'code' => 0,
            'msg' => 'ok'
        ]));
    }

    /**
     * @return Response
     */
    public function settings(): Response
    {
        $model = Common::getModelInfo('midjourney');
        $settings = $model->settings;
        $userId = session('user.id') ?? session('user.uid');
        return json([
            'code' => 0,
            'msg' => 'ok',
            'data' => [
                'allowAllFast' => ($settings['allowAllFast']['value'] ?? false) || User::isVip($userId),
                'hideSpeedOptions' => $settings['hideSpeedOptions']['value'] ?? false,
            ]
        ]);
    }

    /**
     * 格式化提示词
     * @param $image
     * @param $text
     * @return string
     */
    protected function formatPrompt($image, $text): string
    {
        $params = "";
        if ($text) {
            $params .= $text;
        }
        if (!empty($image['styles'])) {
            $params .= ". " . implode(", ", $image['styles']);
        }
        $prompt = $image['prompt'] ?? '';
        $prompt = preg_replace(['/--fast/i', '/--relax/'], '', $prompt);
        if (strpos($prompt, '--ar') === false) {
            if (isset($image['ar']) && $image['ar'] !== "customized") {
                if ($image['ar'] !== "1:1") {
                    $params .= " --ar " . $image['ar'];
                }
            } else if (!empty($image['widthRatio']) && !empty($image['heightRatio'])) {
                $params .= " --ar " . $image['widthRatio'] . ":" . $image['heightRatio'];
            }
        }
        if (strpos($prompt, '--c') === false) {
            if (!empty($image['chaos'])) {
                $params .= " --c " . $image['chaos'];
            }
        }
        if (strpos($prompt, '--s') === false) {
            if (isset($image['stylize']) && $image['stylize'] !== 100) {
                $params .= " --s " . $image['stylize'];
            }
        }
        if (strpos($prompt, '--iw') === false) {
            if (!empty($image['refs'])) {
                $params .= " --iw " . $image['iw'];
            }
        }
        if (strpos($prompt, '--zoom') === false) {
            if (isset($image['zoom'])) {
                $params .= " --zoom " . $image['zoom'];
                $image['zoom'] = "";
            }
        }
        if (strpos($prompt, '--niji') === false) {
            if (isset($image['model']) && $image['model'] === "niji") {
                $params .= " --niji";
            }
        }
        if (isset($image['others'])) {
            $params .= " " . $image['others'];
        }
        if (strpos($prompt, '--v') === false) {
            if (isset($image['version']) && intval($image['version']) !== 6) {
                $params .= " --v " . $image['version'];
            }
        }
        if (empty($image['fast'])) {
            $params .= ' --relax';
        }
        return $params;
    }

}